#!/usr/bin/env python
#
#       GeneNetWalk_interface.py
#       
#       Copyright 2009 Brandon Invergo <brandon@brandoninvergo.com>
#       
#       This program is free software; you can redistribute it and/or modify
#       it under the terms of the GNU General Public License as published by
#       the Free Software Foundation; either version 2 of the License, or
#       (at your option) any later version.
#       
#       This program is distributed in the hope that it will be useful,
#       but WITHOUT ANY WARRANTY; without even the implied warranty of
#       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#       GNU General Public License for more details.
#       
#       You should have received a copy of the GNU General Public License
#       along with this program; if not, write to the Free Software
#       Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#       MA 02110-1301, USA.

import sys
import locale
import platform

try:
    import pygtk
    #tell pyGTK, if possible, that we want GTKv2
    pygtk.require("2.0")
except:
    #Some distributions come with GTK2, but not pyGTK
    pass

try:
    import gtk
    import gtk.glade
    
except:
    print "You need to install pyGTK or GTKv2 ",
    print "or set your PYTHONPATH correctly."
    print "try: export PYTHONPATH=",
    print "/usr/local/lib/python2.2/site-packages/"
    sys.exit(1)
    
import pango
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Seq import MutableSeq
from Bio.Seq import Seq
from Bio.Alphabet import IUPAC
from Bio.Data import CodonTable
from StringIO import StringIO
import random
from math import log

class GeneNetWalk_GUI:
    # Get the local text encoding and define a function to convert to
    # UTF-8 format
    encoding = locale.getdefaultlocale()[1]
    utf8conv = lambda self, x : unicode(x, self.encoding).encode('utf8')
    
    # A list of sequences that have been entered
    input_sequences = []
    
    # A list of sequences that have been created by mutation
    output_proteins = []
    
    # The clipboard for copy/paste
    clipboard = gtk.clipboard_get(gtk.gdk.SELECTION_CLIPBOARD)
    
    # Application preferences:
    # translation table
    trans_table = 11
    # allow internal stop codons
    allow_stops = False
    # which mutations: "nonsyn"(synonymous), "syn"(onymous), "both"
    which_mutations = "nonsyn"
    
    # This function reads the sequences from the Input text area and stores
    # them in the input_sequences list
    def read_sequences(self, data):
        # Create a handle of the text data
        IO_handle = StringIO(data)
        # Parse the text data in the FASTA format
        all_records = list(SeqIO.parse(IO_handle, "fasta"))
        
        # If no sequences were parsed out, then the format was wrong (or
        # just empty)
        if len(all_records) == 0:
            self.error.format_secondary_text("Input data is not in a proper FASTA format.")
            self.error.show()
            return
        
        # For each sequence record parsed out, set the alphabet to a DNA
        # alphabet. If the sequence is in lower case letters, convert to
        # upper case. Finally, append to the sequence list
        for record in all_records:
            record.seq.alphabet = IUPAC.unambiguous_dna
            if str(record.seq).islower():
                record.seq = Seq(str(record.seq).upper(), IUPAC.unambiguous_dna)
            self.input_sequences.append(record)
                    
    # This function does point mutations until a non-synonymous sequence
    # is found
    def do_mutation(self, record, n):
        # Get the sequence length
        sequence_length = len(record.seq)
        # Get the protein product of the original sequence
        try:
            orig_protein = record.seq.translate(table=self.trans_table)
        except CodonTable.TranslationError, TransError:
            rec_name = record.name
            self.error.format_secondary_text(str(TransError) + " (in original sequence: %s)"%rec_name)
            self.error.show()
            return
        
        if orig_protein.find("*") != len(orig_protein)-1 and self.allow_stops == False:
            rec_name = record.name
            self.error.format_secondary_text("Internal stop codon in original sequence: %s)"%rec_name)
            self.error.show()
            return
        
        # While we still haven't found a non-synonymous sequence
        stop = False
        while stop == False:
            # Define the available nucleotides for a mutation
            nucleotide_list = ['A', 'C', 'T', 'G']
            # Create a mutable copy of the sequence
            mutable_seq = record.seq.tomutable()
            # Choose a random position in the sequence, skipping the first
            # and last 3 positions, which compose the start and stop codons.
            rand_pos = random.randint(3, sequence_length-4)
            # Identify the current nucleotide in that position
            nucleotide = record.seq[rand_pos]
            # Remove the current nucleotide from the list of available 
            # nucleotides
            nucleotide_list.remove(nucleotide)
            # Pick a random nucleotide from the available ones
            new_nucleotide = random.choice(nucleotide_list)
            # Change the nucleotide in the mutable sequence to the new
            # one
            mutable_seq[rand_pos] = new_nucleotide
            # Get the protein product of the new sequence
            try:
                mutant_protein = mutable_seq.toseq().translate(table=self.trans_table)
            except CodonTable.TranslationError, TransError:
                pass     
            # If the new product is different than the original and the
            # mutation did not introduce a stop codon, break out of the loop
            if (self.which_mutations == "nonsyn" and str(mutant_protein) != str(orig_protein)) \
                or (self.which_mutations == "syn" and str(mutant_protein) == str(orig_protein) \
                or (self.which_mutations == "both")):
                if self.allow_stops or mutant_protein.find("*") == len(mutant_protein)-1:
                    if self.output_proteins.count(str(mutant_protein)) == 0:
                        self.output_proteins.append(str(mutant_protein))
                        stop = True
        
        # Create a new mutant sequence record. The sequence is the mutant
        # created in the while loop. The id is the original id with a 
        # tag appended to designate the mutant number; ie the first mutant
        # generated would have ".m1", the second ".m2" etc appended
        mutant_record = SeqRecord(mutable_seq.toseq(),\
            id = record.id+".m%i"%n, \
            description = record.description)
        # Return the mutant sequence, the mutated codon and and the protein seq
        return (mutant_record, rand_pos, mutant_protein)
        
    # Calculate summary statistics
    def do_summary(self, record):
        seq_length = len(record)
        
        output_iter = self.summary_textbuffer.get_end_iter()
        
        self.summary_textbuffer.insert(output_iter, ">%s\n"%record.id)
        
        self.summary_textbuffer.insert(output_iter, "Length: %s bp\n"%seq_length)
        
        # 3 possible mutations at each base pair
        num_mutants = 3*seq_length
        self.summary_textbuffer.insert(output_iter, "Total possible mutations: %i\n"%num_mutants)
        
        num_syn = 0
        num_non = 0
        rec_seq = record.seq
        
        # For each codon, except the start and stop codons, try each possible
        # mutation and check to see if it's synonymous on non-synonymous
        # User preferences are obeyed
        for pos in range(3, seq_length-3, 3):
            temp_seq = rec_seq[pos:pos+3]
            amino = temp_seq.translate(table=self.trans_table)
            for base in range(0,3):
                mutable_seq = temp_seq.tomutable()
                nucleotide_list = ['A', 'C', 'T', 'G']
                nucleotide = temp_seq[base]
                nucleotide_list.remove(nucleotide)
                for nuc in nucleotide_list:
                    mutable_seq[base] = nuc
                    temp_amino = mutable_seq.toseq().translate(table=self.trans_table)
                    if str(temp_amino) == "*" and self.allow_stops == False:
                        continue
                    if str(temp_amino) == str(amino):
                        num_syn = num_syn + 1
                    else:
                        num_non = num_non + 1
                        
        total_valid_muts = num_syn + num_non
        
        if self.allow_stops == False:
            self.summary_textbuffer.insert(output_iter, "Total valid mutations: %i "%total_valid_muts)
            self.summary_textbuffer.insert(output_iter, "(excluding mutations that produce in-frame stop codons and mutations in the gene's start and stop codons)\n")
            self.summary_textbuffer.insert(output_iter, "Total valid synonymous mutations: %i\n"%num_syn)
            self.summary_textbuffer.insert(output_iter, "Total valid non-synonymous mutations: %i\n"%num_non)
        else:
            self.summary_textbuffer.insert(output_iter, "Total valid mutations: %i "%total_valid_muts)
            self.summary_textbuffer.insert(output_iter, "(excluding mutations in the frame's start and stop codons)\n")
            self.summary_textbuffer.insert(output_iter, "Total synonymous mutations: %i\n"%num_syn)
            self.summary_textbuffer.insert(output_iter, "Total non-synonymous mutations: %i\n"%num_non)
            
        self.summary_textbuffer.insert(output_iter, "\n")
        
        return total_valid_muts
        
    # This function defines what happens when the user clicks the "Generate"
    # button
    def on_generate_button(self, widget):
        # Make sure that the list of sequences is empty before creating
        # new ones
        self.input_sequences = []
        
        # Ditto for the output proteins
        self.output_proteins = []
        
        # Clear the current output
        self.dna_textbuffer.set_text("")
        self.protein_textbuffer.set_text("")
        self.summary_textbuffer.set_text("")
        
        # Get all of the text in the Input text area
        start_iter = self.input_textbuffer.get_start_iter()
        end_iter = self.input_textbuffer.get_end_iter()
        input_text = self.input_textbuffer.get_text(start_iter, end_iter)
        
        # Get the number of mutants to create
        num_mutants = int(self.num_mutants_entry.get_text())
        
        # Read in the sequences
        self.read_sequences(input_text)
        
        # For each sequence record...
        for record in self.input_sequences:
            
            total_valid_muts = self.do_summary(record)
            if total_valid_muts < num_mutants:
                self.error.format_secondary_text("# of mutants to generate is larger than the total valid mutations for %s (%i)"%(record.id, total_valid_muts))
                self.error.show()
                continue
            
            # Create n mutants
            for n in range(1, num_mutants+1):
                # Create a mutant
                mutant = self.do_mutation(record, n)
                if mutant == None:
                    continue
                
                # Print the mutant's DNA sequence in the Output text area
                output_iter = self.dna_textbuffer.get_end_iter()
                start_line = output_iter.get_line()
                mutant_string = mutant[0].format("fasta")
                self.dna_textbuffer.insert(output_iter, mutant_string + '\n')
                
                # Find the line and position in the printout where the
                # mutation is. Biopython formats FASTA sequences to have
                # 60 characters per line.
                mutation_line = int(mutant[1] / 60)
                mutation_char = mutant[1] % 60
                
                # Highlight the position where the mutation occured
                highlight_start_iter = self.dna_textbuffer.get_iter_at_line(start_line)
                highlight_start_iter.forward_lines(mutation_line + 1)
                highlight_start_iter.forward_chars(mutation_char)
                highlight_end_iter = highlight_start_iter.copy()
                highlight_end_iter.forward_char()
                self.dna_textbuffer.apply_tag_by_name("highlight",\
                    highlight_start_iter, highlight_end_iter)
                    
                # Print the mutant's protein sequence in the output text area
                output_iter = self.protein_textbuffer.get_end_iter()
                start_line = output_iter.get_line()
                self.protein_textbuffer.insert(output_iter, ">%s\n"%mutant[0].id)
                self.protein_textbuffer.insert(output_iter, str(mutant[2]) + "\n\n")
                highlight_start_iter = output_iter.copy()
                highlight_start_iter.backward_lines(2)
                #self.protein_textbuffer.apply_tag_by_name("highlight",\
                #    highlight_start_iter, output_iter)
                amino_loc = int(mutant[1]/3)
                highlight_start_iter.forward_chars(amino_loc)
                highlight_end_iter = highlight_start_iter.copy()
                highlight_end_iter.forward_char()
                self.protein_textbuffer.apply_tag_by_name("highlight",\
                    highlight_start_iter, highlight_end_iter)
    
    def on_menu_about(self, widget):
        self.about_dialog.show()
        
    def on_menu_cut(self, widget):
        focus = self.main_window.get_focus()
        if focus == self.input_textview:
            self.input_textbuffer.cut_clipboard(self.clipboard, True)
        elif focus == self.dna_textview:
            self.dna_textbuffer.copy_clipboard(self.clipboard)
        elif focus == self.protein_textview:
            self.protein_textbuffer.copy_clipboard(self.clipboard)
        elif focus == self.summary_textview:
            self.summary_textbuffer.copy_clipboard(self.clipboard)
    
    def on_menu_copy(self, widget):
        focus = self.main_window.get_focus()
        if focus == self.input_textview:
            self.input_textbuffer.copy_clipboard(self.clipboard)
        elif focus == self.dna_textview:
            self.dna_textbuffer.copy_clipboard(self.clipboard)
        elif focus == self.protein_textview:
            self.protein_textbuffer.copy_clipboard(self.clipboard)
        elif focus == self.summary_textview:
            self.summary_textbuffer.copy_clipboard(self.clipboard)
        
    def on_menu_paste(self, widget):
        focus = self.main_window.get_focus()
        if focus == self.input_textview:
            self.input_textbuffer.paste_clipboard(self.clipboard, None, True)
    
    def on_menu_import(self, widget):
        self.import_dialog.show()
        
    def on_import_cancel_button(self, widget):
        self.import_dialog.hide()
        
    def on_import_open_button(self, widget):
        filename = self.import_dialog.get_filename()
        self.import_file(filename)
        
    def on_import_file_activated(self, widget):
        filename = self.import_dialog.get_filename()
        self.import_file(filename)
        
    def import_file(self, filename):
        try:
            opened_file = open(filename, "r")
        except:
            self.error.format_secondary_text("Error opening file.")
            self.error.show()
        else:
            content = opened_file.read()
            #if platform.system() == 'Linux':
            try:
                content = self.utf8conv(content)  
            except:
                self.error.format_secondary_text("Unsupported file type.")
                self.error.show()
                return
            self.input_textbuffer.set_text(content)
            self.dna_textbuffer.set_text("")
            self.protein_textbuffer.set_text("")
            self.summary_textbuffer.set_text("")
            opened_file.close()
            self.import_dialog.hide()
        
    def on_menu_export(self, widget):
        self.export_dialog.show()
        
    def on_export_cancel_button(self, widget):
        self.export_dialog.hide()
        
    def on_export_save_button(self, widget):
        filename = self.export_dialog.get_filename()
        self.export_file(filename)
        
    def on_export_file_activated(self, widget):
        filename = self.export_dialog.get_filename()
        self.export_file(filename)
        
    def export_file(self, filename):
        start_iter = self.dna_textbuffer.get_start_iter()
        end_iter = self.dna_textbuffer.get_end_iter()
        output_text = self.dna_textbuffer.get_text(start_iter, end_iter)
        if filename.rpartition(".")[2] != "fasta":
            filename = filename + ".fasta"
        try:
            new_file = open(filename, "w")
        except:
            self.error.set_secondary_text("Error opening file for saving.")
            self.error.show()
            return
        try:
            new_file.write(output_text)
        except:
            self.error.set_secondary_text("Error writing file.")
            self.error.show()
        else:
            self.export_dialog.hide()      
        
    def on_error_response(self, widget, response):
        self.error.hide()
        
    def on_about_dialog_response(self, widget, response):
        self.about_dialog.hide()
        
    def on_menu_prefs(self, widget):
        trans_table_spin = self.glade_gui.get_widget("prefs_trans_table")
        trans_table_spin.set_value(self.trans_table)
                
        stop_codons_check = self.glade_gui.get_widget("prefs_stop_codons")
        stop_codons_check.set_active(self.allow_stops)
        
        mut_radio_group = self.glade_gui.get_widget("prefs_nonsyn_radio").get_group()
        for radio in mut_radio_group:
            if radio.get_name() == "prefs_%s_radio"%self.which_mutations:
                radio.set_active(True)
                break
        
        res = self.prefs_dialog.run()
        if res == gtk.RESPONSE_CLOSE:
            self.trans_table = trans_table_spin.get_value_as_int()
            
            
            self.allow_stops = stop_codons_check.get_active()
            
            
            for radio in mut_radio_group:
                if radio.get_active():
                    name = radio.get_name().split("_")
                    self.which_mutations = name[1]
                    break
                    
        self.prefs_dialog.hide()
        
    def close_application(self, widget):
        gtk.main_quit()
    
    def __init__(self):
        gladefile="GeneNetWalk.glade"
        self.glade_gui = gtk.glade.XML(gladefile)
        
        callbacks = {"gtk_main_quit" : self.close_application, \
            "on_about_dialog_response" : self.on_about_dialog_response,\
            "on_menu_import" : self.on_menu_import, \
            "on_menu_export" : self.on_menu_export, \
            "on_menu_about"  : self.on_menu_about, \
            "on_menu_cut" : self.on_menu_cut,\
            "on_menu_copy": self.on_menu_copy,\
            "on_menu_paste": self.on_menu_paste,\
            "on_generate_button" : self.on_generate_button,\
            "on_import_cancel_button" : self.on_import_cancel_button,\
            "on_import_open_button" : self.on_import_open_button,\
            "on_import_file_activated": self.on_import_file_activated,\
            "on_export_cancel_button" : self.on_export_cancel_button,\
            "on_export_save_button" : self.on_export_save_button,\
            "on_export_file_activated": self.on_export_file_activated,\
            "on_error_dialog_response": self.on_error_response,\
            "on_menu_prefs" : self.on_menu_prefs}
        
        self.glade_gui.signal_autoconnect(callbacks)
        
        self.main_window = self.glade_gui.get_widget("main_window")
        self.main_window.connect("destroy", self.close_application)
        self.import_dialog = self.glade_gui.get_widget("import_file_dialog")
        self.export_dialog = self.glade_gui.get_widget("export_file_dialog")
        self.about_dialog = self.glade_gui.get_widget("about_dialog")
        self.error = self.glade_gui.get_widget("error_dialog")
        self.input_textview = self.glade_gui.get_widget("input_textview")
        self.dna_textview = self.glade_gui.get_widget("dna_textview")
        self.protein_textview = self.glade_gui.get_widget("protein_textview")
        self.summary_textview = self.glade_gui.get_widget("summary_textview")
        self.num_mutants_entry = self.glade_gui.get_widget("num_mutants_entry")
        self.prefs_dialog = self.glade_gui.get_widget("prefs_dialog")
        
        self.input_textbuffer = gtk.TextBuffer()
        self.input_textview.set_buffer(self.input_textbuffer)
        self.input_textview.modify_font(pango.FontDescription("monospace"))
   
        self.dna_textbuffer = gtk.TextBuffer()
        self.dna_textview.set_buffer(self.dna_textbuffer)
        self.dna_textbuffer.create_tag("highlight", background = "red")
        self.dna_textbuffer.create_tag("font", font = "fixed")
        self.dna_textview.modify_font(pango.FontDescription("monospace"))
        
        self.protein_textbuffer = gtk.TextBuffer()
        self.protein_textview.set_buffer(self.protein_textbuffer)
        self.protein_textbuffer.create_tag("highlight", background = "red")
        self.protein_textbuffer.create_tag("font", font = "fixed")
        self.protein_textview.modify_font(pango.FontDescription("monospace"))
        
        self.summary_textbuffer = gtk.TextBuffer()
        self.summary_textview.set_buffer(self.summary_textbuffer)
        self.summary_textbuffer.create_tag("font", font = "fixed")
        self.summary_textview.modify_font(pango.FontDescription("monospace"))
        
        self.prefs_dialog.add_button(gtk.STOCK_CLOSE, gtk.RESPONSE_CLOSE)
        

def main():
    """The main function"""
    # Enter the main Gtk loop
    gtk.main()
    return 0
        
if __name__ == "__main__":
    GeneNetWalk_GUI()
    main()  
